

import fiftyone.utils.data as foud
import fiftyone.core.metadata as fom
import fiftyone.core.labels as fol
import os

class YOLOv8LabeledImageDatasetImporter(foud.LabeledImageDatasetImporter):
    """Custom importer for labeled image datasets.

    Args:
        dataset_dir (None): the dataset directory. This may be optional for
            some importers
        shuffle (False): whether to randomly shuffle the order in which the
            samples are imported
        seed (None): a random seed to use when shuffling
        max_samples (None): a maximum number of samples to import. By default,
            all samples are imported
        **kwargs: additional keyword arguments for your importer
    """

    def __init__(
        self,
        dataset_dir=None,
        shuffle=False,
        seed=None,
        max_samples=None,
        image_folder_name = "Images",
        label_file_name = "Labels",
        field_name = "gt",
        **kwargs,
    ):
        super().__init__(
            dataset_dir=dataset_dir,
            shuffle=shuffle,
            seed=seed,
            max_samples=max_samples,
        )
        self.image_folder_name = image_folder_name
        self.label_file_name = label_file_name
        self.field_name = field_name

    def __len__(self):
        """The total number of samples that will be imported.

        Raises:
            TypeError: if the total number is not known
        """
        # Return the total number of samples in the dataset (if known)
        return len(self.samples)

    def __next__(self):
        """Returns information about the next sample in the dataset.

        Returns:
            an  ``(image_path, image_metadata, label)`` tuple, where

            -   ``image_path``: the path to the image on disk
            -   ``image_metadata``: an
                :class:`fiftyone.core.metadata.ImageMetadata` instances for the
                image, or ``None`` if :meth:`has_image_metadata` is ``False``
            -   ``label``: an instance of :meth:`label_cls`, or a dictionary
                mapping field names to :class:`fiftyone.core.labels.Label`
                instances, or ``None`` if the sample is unlabeled

        Raises:
            StopIteration: if there are no more samples to import
        """
        # Implement loading the next sample in your dataset here
        image_path, image_metadata, detections = next(self.image_iter)
        # print(image_metadata)
        return image_path, image_metadata, {self.field_name: detections}

    @property
    def has_dataset_info(self):
        """Whether this importer produces a dataset info dictionary."""
        return True

    @property
    def has_image_metadata(self):
        """Whether this importer produces
        :class:`fiftyone.core.metadata.ImageMetadata` instances for each image.
        """
        return True

    @property
    def label_cls(self):
        """The :class:`fiftyone.core.labels.Label` class(es) returned by this
        importer.

        This can be any of the following:

        -   a :class:`fiftyone.core.labels.Label` class. In this case, the
            importer is guaranteed to return labels of this type
        -   a list or tuple of :class:`fiftyone.core.labels.Label` classes. In
            this case, the importer can produce a single label field of any of
            these types
        -   a dict mapping keys to :class:`fiftyone.core.labels.Label` classes.
            In this case, the importer will return label dictionaries with keys
            and value-types specified by this dictionary. Not all keys need be
            present in the imported labels
        -   ``None``. In this case, the importer makes no guarantees about the
            labels that it may return
        """
        # Return the appropriate value here
        return {self.field_name: fol.Detections}

    def _parse_label(self, label_path):
        with open(label_path, 'r') as f:
            lines = f.readlines()
        detections = []
        for line in lines:
            line = line.strip()
            if not line:
                continue
            label, cx, cy, w, h = line.split()
            x, y, w, h = float(cx)-float(w)/2, float(cy)-float(h)/2, float(w), float(h)
            detections.append(fol.Detection(label=label, bounding_box=[x, y, w, h]))
        # if not detections:
        #     return None
        return fol.Detections(detections=detections)
    
    def setup(self):
        """Performs any necessary setup before importing the first sample in
        the dataset.

        This method is called when the importer's context manager interface is
        entered, :func:`DatasetImporter.__enter__`.
        """
        self.images_dir = os.path.join(self.dataset_dir, self.image_folder_name)
        self.labels_dir = os.path.join(self.dataset_dir, self.label_file_name)
        if not os.path.isdir(self.images_dir) or not os.path.isdir(self.labels_dir):
            raise Exception(f"Invalid dataset directory {self.images_dir}")

        self.image_paths = sorted([f for f in os.listdir(self.images_dir) if f.endswith('.jpg') or f.endswith('.png') ])
        self.samples = []
        for f in self.image_paths:
            img_path = os.path.join(self.images_dir, f)
            image_metadata = fom.ImageMetadata.build_for(img_path)
            label_path = os.path.join(self.labels_dir, f.replace('.jpg', '.txt'))
            if not os.path.isfile(label_path):
                print(f"Warning: {label_path} not founded")
                detections = fol.Detections(detections=[])
            else:
                detections = self._parse_label(label_path)
            # print(detections)
            self.samples.append((img_path, image_metadata, detections))
            
            
        self.image_iter = iter(self.samples)
        
        self.info = {"dataset_dir": self.dataset_dir, "default_classes": [ "defect_"+str(i) for i in range(10) ]}

    def get_dataset_info(self):
        """Returns the dataset info for the dataset.

        By convention, this method should be called after all samples in the
        dataset have been imported.

        Returns:
            a dict of dataset info
        """
        return self.info

import fiftyone.types as fot

class YOLOv8ImageDataset(fot.LabeledImageDataset):
    def get_dataset_importer_cls(self):
        return YOLOv8LabeledImageDatasetImporter
    def get_dataset_exporter_cls(self):
        pass
